# Embeddings_models.py
# This file contains the Pydantic models for the Embeddings API.
#
# Imports
from typing import List, Literal, Union, Optional
#
# Third-Party Libraries
from pydantic import Field, BaseModel, ConfigDict
#
# Local Libraries
#
#######################################################################################################################
#
# --- Pydantic Models (OpenAI Spec Compliant) ---

class CreateEmbeddingRequest(BaseModel):
    model_config = ConfigDict(
        extra='forbid',
        json_schema_extra={
            "example": {
                "model": "text-embedding-ada-002",
                "input": "The food was delicious and the waiter was friendly.",
                "encoding_format": "float"
            }
        }
    )

    input: Union[str, List[str], List[int], List[List[int]]] = Field(
        ...,
        description="Input text to embed, encoded as a string or array of tokens. To embed multiple inputs in a single request, pass an array of strings or array of token arrays. The input must not exceed the max input tokens for the model (8192 tokens for `text-embedding-ada-002`), cannot be an empty string, and any array must be 2048 dimensions or less.",
        examples=["The food was delicious and the waiter..."]
    )
    model: str = Field( # Keep this as str, the library handles the default if not provided.
        # default=None, # Making model optional here would deviate from OpenAI spec, which requires it.
                         # The library can still use its internal default if the passed model IS the default.
        description="ID of the model to use. You can use the [List models](/docs/api-reference/models/list) API to see all of your available models, or see our [Model overview](/docs/models) for descriptions of them.",
        examples=["text-embedding-ada-002"]
    )
    encoding_format: Optional[Literal["float", "base64"]] = Field(
        default="float",
        description="The format to return the embeddings in. Can be either `float` or `base64`."
    )
    dimensions: Optional[int] = Field(
        default=None,
        description="The number of dimensions the resulting output embeddings should have. Only supported in `text-embedding-3` and later models.",
        ge=1
    )
    user: Optional[str] = Field(
        default=None,
        description="A unique identifier representing your end-user, which can help OpenAI to monitor and detect abuse."
    )

class EmbeddingData(BaseModel):
    object: Literal["embedding"] = "embedding"
    embedding: Union[List[float], str] = Field(description="The embedding vector, which is a list of floats or a base64-encoded string. The length of vector depends on the model.")
    index: int = Field(description="The index of the embedding in the list of embeddings.")

class EmbeddingUsage(BaseModel):
    prompt_tokens: int = Field(description="The number of tokens used by the prompt.")
    total_tokens: int = Field(description="The total number of tokens used by the request.")

class CreateEmbeddingResponse(BaseModel):
    object: Literal["list"] = "list"
    data: List[EmbeddingData] = Field(description="The list of embeddings generated by the model.")
    model: str = Field(description="The model ID used to generate the embeddings.")
    usage: EmbeddingUsage = Field(description="The usage information for the request.")

#
# End of Pydantic Models
#######################################################################################################################
